#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import subprocess

# --- 配置区 ---

# Beeline 连接字符串
BEELINE_CONNECT_URL = "jdbc:hive2://hadoop102:10000"
# Hive 用户
HIVE_USER = "atguigu"
# Hive 数据库名
HIVE_DB = "car_analysis_db"

# 要导出的 Hive 表列表
TABLES_TO_EXPORT = [
    "level_sales_relation",
    "level_price_relation",
    "manufacturer_rank"
]

# 本地存放结果的目录
LOCAL_RESULTS_DIR = "/home/atguigu/spark_analysis/data/data_result"


# --- 函数区 ---

def execute_beeline_query_to_csv(query, output_filepath):
    """
    使用 Beeline 执行 Hive 查询并将结果保存为 CSV 文件。
    """
    command = (
        f"beeline -u {BEELINE_CONNECT_URL} -n {HIVE_USER} "
        f"--outputformat=csv2 --showDbInPrompt=false "
        f"-e \"{query}\" > \"{output_filepath}\""
    )

    print(f"[*] 准备执行命令: {command}")

    try:
        # **修改点**: 兼容 Python 3.6
        # 1. 移除 text=True 参数
        # 2. 添加 universal_newlines=True 来自动进行文本解码
        result = subprocess.run(
            command,
            shell=True,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True  # 替代 text=True
        )
        print(f"[SUCCESS] 查询结果已成功保存到: {output_filepath}")

        if result.stderr:
            print("[INFO] Beeline 执行日志:")
            print(result.stderr.strip())

        return True
    except subprocess.CalledProcessError as e:
        print(f"[ERROR] 执行 Beeline 命令失败!")
        print(f"  - 返回码: {e.returncode}")
        print(f"  - 错误输出: {e.stderr.strip()}")
        return False
    except FileNotFoundError:
        print("[ERROR] 'beeline' 命令未找到。请确保它在系统的PATH中。")
        return False


def main():
    """主执行函数"""
    print("--- 开始执行 Hive 数据导出脚本 ---")

    try:
        os.makedirs(LOCAL_RESULTS_DIR, exist_ok=True)
        print(f"[INFO] 已确保本地结果目录存在: {LOCAL_RESULTS_DIR}")
    except OSError as e:
        print(f"[ERROR] 创建本地目录失败: {LOCAL_RESULTS_DIR}")
        print(f"  - 错误原因: {e}")
        return

    print("\n--- 开始导出 Hive 表 ---")
    for table_name in TABLES_TO_EXPORT:
        print(f"\n--- 正在处理表: {table_name} ---")

        query_sql = f"SELECT * FROM {HIVE_DB}.{table_name};"
        output_csv_path = os.path.join(LOCAL_RESULTS_DIR, f"{table_name}.csv")

        execute_beeline_query_to_csv(query_sql, output_csv_path)

    print("\n--- 脚本执行完毕 ---")
    print(f"所有分析结果已导出到目录: {LOCAL_RESULTS_DIR}")


if __name__ == "__main__":
    main()